import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 5)))
sys.path.insert(0, parent_path)

from base64 import b64encode
import traceback
import numpy as np
import socket, json, cv2, yaml, time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.PIL_draw import pic_text
from core.utils.visualize import get_color_map_list
from core.algorithm.deeplearning.paddle.detectors.seg.seg_utils import visualize as seg_vis


class FTPProcessor(BaseProcessor):
    """FTP检测后处理方案
    """
    def __init__(self, 
                 keys,
                 Custom_cls_map,
                 Custom_target_cls,
                 Custom_pixel_thresholds,
                 ) -> None:
        self.keys = keys
        if isinstance(Custom_target_cls,str):
            with open(Custom_target_cls,"r") as f:
                self.target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.target_cls=set(Custom_target_cls)
        if isinstance(Custom_cls_map,str):
            with open(Custom_cls_map,"r") as f:
                self.cls_map = {x:y for x,y in enumerate(yaml.safe_load(f)["label_list"])}
        else:
            self.cls_map = Custom_cls_map
        self.pixel_thresholds = Custom_pixel_thresholds
        self.color_list = get_color_map_list(len(self.cls_map))
        
    def init_check(self):
        assert "seg" in self.keys
        super().init_check()
    
    def __call__(self, data:Dict) -> Dict:
        
        batch_img = data[self.keys["in"]].copy()
        cover_result = data[self.keys["seg"]]
        objects = {}
        # start_idx = 0
        for idx_, mask in enumerate(cover_result.get("masks",[])):
            download_path = data["data"][idx_]
            objects[download_path]={}
            for cls_id, cls_name in self.cls_map.items():
                mask:np.ndarray
                if not cls_id in self.target_cls: # 只有在配置文件中指定的类才进行判断
                    continue
                objects[download_path][cls_name] = False
                pixel_threshold = self.pixel_thresholds[cls_id]*mask.shape[0]*mask.shape[1]
                if np.sum(mask[mask == cls_id])>pixel_threshold:
                    objects[download_path][cls_name] = True
            objects[download_path]["img_size"] = [batch_img[idx_].shape[1], batch_img[idx_].shape[0]]
        vis_img = seg_vis(batch_img,cover_result["masks"])
        data[self.keys["out_vis"]] = vis_img
        data[self.keys["out_obj"]] = objects
            
        return data

    def close(self):
        self.status = 0